package com.ruoyi.bearing.controller;

import com.ruoyi.bearing.domain.ModelTrain;
import com.ruoyi.bearing.mapper.ModelTrainMapper;
import com.ruoyi.bearing.service.IModelTrainService;
import com.ruoyi.common.core.controller.BaseController;
import com.ruoyi.common.core.domain.AjaxResult;
import com.ruoyi.common.exception.base.BaseException;
import com.ruoyi.common.exception.flask.FlaskException;
import io.swagger.annotations.Api;
import io.swagger.annotations.ApiOperation;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.security.access.prepost.PreAuthorize;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.multipart.MultipartFile;

import java.io.IOException;

/**
 * 模型训练
 */
@Api("模型训练")
@RestController
@RequestMapping("/model/train")
public class ModelTrainController extends BaseController {
    @Autowired
    private IModelTrainService iModelTrainService;
    @Autowired
    ModelTrainMapper modelTrainMapper;

    /**
     * 诊断模型训练
     */
    @ApiOperation("故障诊断模型训练")
    @PreAuthorize("@ss.hasPermi('model:train:setDiagParam')")
    @PostMapping("/setDiagParam")
    @Transactional(rollbackFor = Exception.class)
    public AjaxResult setDiagParam(@RequestParam("files") MultipartFile[] files, ModelTrain modelTrain) throws IOException, IllegalAccessException {
        if (modelTrainMapper.selectModelTrainByModelName(modelTrain.getName()) != null) {
            throw new BaseException("模型名称已存在");
        }
        String res = iModelTrainService.startDiagnosisModelTrain(files[0], modelTrain);
        if ("".equals(res)) {
            throw new FlaskException();
        }
        if (res.contains("请检查文件格式")) {
            throw new BaseException("请检查文件格式");
        }
        return AjaxResult.success(res);
    }

    /**
     * 预测模型训练
     */
    @ApiOperation("预测模型训练")
    @PreAuthorize("@ss.hasPermi('model:train:setRulParam')")
    @PostMapping("/setRulParam")
    @Transactional(rollbackFor = Exception.class)
    public AjaxResult setRulParam(@RequestParam("files") MultipartFile[] files, ModelTrain modelTrain) throws IOException, IllegalAccessException {
        if (modelTrainMapper.selectModelTrainByModelName(modelTrain.getName()) != null) {
            throw new BaseException("模型名称已存在");
        }
        String res = iModelTrainService.startPredictionModelTrain(files[0], modelTrain);
        if ("".equals(res)) {
            throw new FlaskException();
        }
        return AjaxResult.success(res);
    }
}
